import gradio as gr
from PIL import Image, ImageOps, ImageChops
from utils import generate_random_dir, generate_random_file
from utils.ps_canvas import PSTools
from hair_change_color import HairChangeColor
import os
from demo_models import hair_change_color_model



def main(origin_image: Image, color: str, b_strength, g_strength, r_strength) -> Image:
    # 设置数据缓存目录
    _save_dir = generate_random_dir()
    # 数据先保存
    origin_filename = generate_random_file(dir_=_save_dir)
    origin_image.save(origin_filename)
    # 调模型
    try:
        model = hair_change_color_model
        filename = model.change_color(origin_filename, _save_dir,
                                      color, b_strength, g_strength, r_strength)
        image = Image.open(filename)
    except Exception as e:
        logger.exception(str(e))
        image = None
    os.system(f"rm -rf {_save_dir}")
    return image


def demo():
    with gr.Row():
        origin_image = gr.Image(
            type="pil", label="原始图像")
        with gr.Column(scale=1):
            color = gr.ColorPicker(label="头发颜色")
            color_strength = gr.Dropdown(
                choices=[1, 2, 3, 4, 5, 8, 10], label="染发强度", value=2)
            # b_strength = gr.Dropdown(choices=[1,2,4,5,8,10],label="蓝色通道染发强度",value=2)
            # g_strength = gr.Dropdown(choices=[1,2,4,5,8,10],label="绿色通道染发强度",value=2)
            # r_strength = gr.Dropdown(choices=[1,2,4,5,8,10],label="红色色通道染发强度",value=2)
        out_image = gr.Image(type="pil")
    gr.Button("运行").click(main, inputs=[
        origin_image, color, color_strength, color_strength, color_strength], outputs=[out_image])
